# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
TF Serving  gRPC 服务调用 demo 代码


# 查看模型版本信息（8501： http端口  lpOCR：模型名称）
http://localhost:8501/v1/models/cifar10
# 查看模型结构信息
http://localhost:8501/v1/models/cifar10/metadata


# 当前项目: Cifar10-Classification
# 开发作者: vincent
# 创建平台: PyCharm Community Edition  TensorFlow 1.12.0
# 版    本: V1.0
"""
import cv2
import numpy as np
import tensorflow as tf
import grpc
from tensorflow_serving.apis import prediction_service_pb2_grpc   # pip install tensorflow_serving-api
from tensorflow_serving.apis import predict_pb2


# gRPC 接口调用程序
class GRPC(object):
    def __init__(self, ip, port, grpc_name, input_name, output_name,
                 time_out=5.0, grpc_signature_name='', http_port=None):
        """
        :param ip:                    IP
        :param port:                  gRPC端口
        :param grpc_name:             模型 名称
        :param input_name:            模型输入名称
        :param output_name:           模型输出名称
        :param time_out:              服务最长等待时间
        :param grpc_signature_name:   Serving名称
        :param http_port:             http 端口
        """
        self.input_name = input_name
        self.output_name = output_name
        self.time_out = time_out
        self.ip = ip
        self.port = port
        self.grpc_name = grpc_name
        self.http_port = http_port

        channel = grpc.insecure_channel('{}:{}'.format(ip, port))
        self.stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
        self.request = predict_pb2.PredictRequest()
        self.request.model_spec.name = grpc_name
        self.request.model_spec.signature_name = grpc_signature_name   # 'serving_default'

        print('Serving: {}:{}  gRPC Name:{}\nInputName:{}  OutputName:{}'.format(
            self.ip, self.port, self.grpc_name, self.input_name, self.output_name))

        # 若提供了 http 端口，可通过http 查看Serving相关信息
        if http_port is not None:
            import requests
            print('{}Serving Information{}'.format('= '*20, ' ='*20))
            # 模型版本信息
            version_info = requests.get('http://{}:{}/v1/models/{}'.format(
                self.ip, self.http_port, self.grpc_name)).text
            print('Model Version Information:\n{}'.format(version_info))
            # 模型输入、输出结构信息
            model_info = requests.get('http://{}:{}/v1/models/{}/metadata'.format(
                self.ip, self.http_port, self.grpc_name)).text
            print('Model Structure Information:\n{}'.format(model_info))
            print('{}'.format('= '*50))

    # Serving 服务调用
    def __call__(self, input_data):
        """
        :param input_data:   输入图片数据，保证与Serving接口要求保持一致
        :return:
        """

        self.request.inputs[self.input_name].CopyFrom(
                tf.contrib.util.make_tensor_proto(input_data))

        try:
            result = self.stub.Predict(self.request, self.time_out)
            return result.outputs[self.output_name].float_val
        except Exception as error:
            print('Error:  {}'.format(error))
            return None

# ##################################################################################
# 创建 gRPC 服务实例
cifar_serving = GRPC(ip='localhost', port='8500', http_port='8501', grpc_name='cifar10', time_out=10,
                     input_name='input_0', output_name='output_0')

# 读取测试图片，按Serving接口要求转换格式
image_path = r'D:\Desktop\Cifar10-Classification\Data\Image\test_data\Test_airplane_00022.bmp'

# 读取图片
image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)
image = cv2.resize(image, (224, 224))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # RGB 通道转换

# 转化为批次数据
data = np.array([image], dtype='float32')

# RPC模型预测
result = cifar_serving(data)
print(('{:.3f}  '*10).format(*result))

